from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
from app.config import MysqlConfig
from app.exceptions.mysql_exception import mysql_url_required_exception
from langchain.prompts import PromptTemplate

PROMPT_SUFFIX = """Only use the following tables:
{table_info}

Question: {input}

"""

_mysql_prompt = """You are a MySQL expert. Given an input question, first create a syntactically correct MySQL query to run, then look at the results of the query and return the answer to the input question.
Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the LIMIT clause as per MySQL. You can order the results to return the most informative data in the database.
Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in backticks (`) to denote them as delimited identifiers.
Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
Pay attention to use CURDATE() function to get the current date, if the question involves "today".

严格遵循以下几点要求:

1.根据用户提出的Question, 结合提供的following tables information, 
自行判断用户的问题和提供的table信息关联度是否符合查询的场景，如果60%不符合table_info描述的场景，则返回空字符串

2.只返回语句，不用给出解释，不用道歉
强调不要写超出 schema 之外的点、边类型

3.根据用户的问题，尝试提炼用户问题中的关键信息并给出对应你认为合理的结果,不需要你去解释和问询
忽略与提问相关的背景信息
你需要返回的内容一定是下面的形式,处理结果相应要求如下,不要出现任何和查询语句无关的内容和任何符号

Use the following format:

Question: Question here
SQLQuery: SQL Query to run
SQLResult: Result of the SQLQuery
Answer: Final answer here

"""

MYSQL_PROMPT = PromptTemplate(
    input_variables=["input", "table_info", 'top_k'],
    template=_mysql_prompt + PROMPT_SUFFIX,
)


class MysqlPromqtManager:

    def __init__(self):
        self.openai_llm = ChatOpenAI(temperature=0, verbose=True, model='gpt-4')
        self.mysql_obj = MysqlConfig()

    @property
    def db(self):
        if not self.mysql_obj.mysql_url:
            raise mysql_url_required_exception()
        return SQLDatabase.from_uri(self.mysql_obj.mysql_url)

    def get_mysql_handler(self, question: str, need_execute=False):
        chain = create_sql_query_chain(llm=self.openai_llm, db=self.db, prompt=MYSQL_PROMPT, k=1000)
        r = QuerySQLDataBaseTool(db=self.db)
        qs = chain | r
        sql_cmd = qs.invoke({"question": question}) if need_execute else chain.invoke({"question": question})
        if sql_cmd.startswith("```sql"):
            sql_cmd = sql_cmd.replace("```sql", "")
        if sql_cmd.endswith("```"):
            sql_cmd = sql_cmd.replace("```", "")
        sql_cmd = sql_cmd.strip()
        sql_cmd_lists = [
            sql_s.strip().replace("\n", " ") for sql_s in sql_cmd.split(";") if
                         sql_s and str(sql_s).startswith("SELECT")]
        return sql_cmd_lists
